import base64
import os
import pymysql
from Crypto.Hash.MD5 import MD5Hash
from erniebot.errors import APIError
from flask import Flask, request, jsonify, abort, Response, send_from_directory
from flask_cors import CORS
from functools import wraps

import cv2
import numpy as np
import erniebot
from pymysql import OperationalError

import dbhelper
import security

HOST: str = "127.0.0.1"
PORT: int = 5000
DB_NAME: str = "EasyEditor_data"
DEBUG: bool = True

app = Flask(__name__)
app.config.from_object(__name__)
CORS(app, resource={r'/*': {'origins': '*'}})


# Configure static file routes
@app.route('/uploads/images/<path:filename>')
def serve_image(filename):
    return send_from_directory('uploads/images', filename)


@app.route('/uploads/docs/<path:filename>')
def serve_doc(filename):
    return send_from_directory('uploads/docs', filename)


def login_required(f):
    @wraps(f)
    def decorated_function(*args, **kwargs):
        token: (str | None) = request.headers.get("Authorization")

        if not token:
            abort(Response("Authorization required.", 401))

        if not token.startswith("Bearer "):
            abort(Response("Unsupported authorization method.", 401))

        jwt_token: str = token.removeprefix("Bearer ")
        jwt_username: (str | None) = security.verify_jwt(jwt_token)

        if not jwt_username:
            abort(Response("Failed to verity the specified token.", 401))

        # Pass username to the view function
        return f(jwt_username, *args, **kwargs)

    return decorated_function


@app.route('/test', methods=['GET'])
def test_availability():
    dbhelper.connect()
    return Response("Backend server is up!", 200)


@app.route('/login', methods=['POST'])
def login():
    username = request.json.get('username')
    password = request.json.get('password')

    if not username or not password:
        abort(Response("Username and password are required.", 400))

    if security.verify_password(username, password):
        token = security.generate_jwt(username)
        return jsonify({
            "success": True,
            "token": token,
            "user": {
                "username": username
            }
        })
    else:
        return jsonify({
            "success": False,
            "message": "Invalid username or password."
        })


@app.route('/add_user', methods=['POST'])
def add_user():
    username = request.form.get("username")
    password = request.form.get("password")

    if not username or not password:
        abort(Response("Username and password are required.", 400))

    try:
        dbhelper.get_user(username)
    except RuntimeError:
        dbhelper.add_user(username, password)
        return jsonify({
            "success": True,
            "message": "User created successfully"
        })
    else:
        return jsonify({
            "success": False,
            "message": "User already exists."
        })


@app.route('/get_docs', methods=['GET'])
@login_required
def get_docs(username: str):
    count, entries = dbhelper.get_documents(username)

    return jsonify({
        "success": True,
        "count": count,
        "documents": entries,
    })


@app.route('/get_doc', methods=['GET'])
@login_required
def get_doc(username: str):
    if not request.args.__contains__("title"):
        return jsonify({
            "success": False,
            "message": "The title of the document is needed.",
        })

    title = request.args["title"]

    try:
        content = dbhelper.get_content(username, title)

        return jsonify({
            "success": True,
            "title": title,
            "content": content,
        })
    except RuntimeError:
        return jsonify({
            "success": False,
            "message": "Cannot find the specified document.",
        })


@app.route('/upload_images', methods=['POST'])
@login_required
def upload_images(username: str):
    img = request.files['file']
    img_name = img.filename
    img_path = f"./uploads/images/{username}/"

    if not os.path.exists(img_path):
        os.makedirs(img_path)

    if img_name is None:
        abort(Response("The name of image is required.", 400))

    file = img.read()
    file_hash = MD5Hash(file).hexdigest()
    file = cv2.imdecode(np.frombuffer(file, np.uint8), cv2.IMREAD_COLOR)

    cv2.imwrite(filename=os.path.join(img_path, img_name), img=file)
    url = f"http://{HOST}:{PORT}/uploads/images/{username}/{img_name}"

    conn = pymysql.connect(host="127.0.0.1", port=3306, user="root", password="root", charset="utf8")
    cursor = conn.cursor()
    conn.commit()

    cursor.execute(f"use {DB_NAME}")
    sql = "insert ignore into `img_path` (`username`, `url`, `img_name`, `img_hash`) values (%s, %s, %s, %s)"
    cursor.execute(sql, (username, url, img_name, file_hash))

    conn.commit()
    cursor.close()
    conn.close()

    result_map = {"url": url}
    return jsonify(result_map)


@app.route('/upload_doc', methods=['POST'])
@login_required
def upload_doc(username: str):
    doc_encoded = request.form["document"]
    doc_path = f"./uploads/docs/{username}/"
    doc_title = request.form["title"]
    doc_filename = f"{MD5Hash(f"{username}_{doc_title}".encode()).hexdigest()}.html"
    doc = base64.b64decode(doc_encoded).decode()

    if not os.path.exists(doc_path):
        os.makedirs(doc_path)

    if doc_title is None:
        abort(Response("The title of the document is required.", 400))

    with open(os.path.join(doc_path, doc_filename), "w") as f:
        f.write(doc)

    url = f"http://{HOST}:{PORT}/uploads/docs/{username}/{doc_filename}"

    conn = pymysql.connect(host="127.0.0.1", port=3306, user="root", password="root", charset="utf8")
    cursor = conn.cursor()
    conn.commit()

    cursor.execute(f"use {DB_NAME}")
    sql = "insert ignore into `doc_path` (`username`, `url`, `file_name`, `title`) values (%s, %s, %s, %s)"
    cursor.execute(sql, (username, url, doc_filename, doc_title))

    conn.commit()
    cursor.close()
    conn.close()

    result_map = {"url": url}
    return jsonify(result_map)


@app.route('/remove_doc', methods=['POST'])
@login_required
def remove_doc(username: str):
    title = request.form["title"]

    if title is None:
        abort(Response("The title of the document is required.", 400))

    try:
        conn = pymysql.connect(host="127.0.0.1", port=3306, user="root", password="root", charset="utf8")
        cursor = conn.cursor()
        conn.commit()

        cursor.execute(f"use {DB_NAME}")
        sql_find = "select `file_name` from `doc_path` WHERE `username` = %s and `title` = %s"

        count: int = cursor.execute(sql_find, (username, title))

        if count == 0:
            return jsonify({
                "success": False,
                "message": "Cannot find the specified document.",
            })

        filename = cursor.fetchone()[0]
        os.remove(f"./uploads/docs/{username}/{filename}")

        sql_delete = "delete from `doc_path` WHERE `username` = %s and `title` = %s"
        cursor.execute(sql_delete, (username, title))

        conn.commit()
        cursor.close()
        conn.close()

        return jsonify({
            "success": True,
        })
    except FileNotFoundError:
        return jsonify({
            "success": True,
        })
    except OperationalError:
        return jsonify({
            "success": False,
        })


@app.route('/get_polish', methods=["POST"])
@login_required
def get_polish(username: str):
    return ai_assist(request.form.get("cont"), "润色")


@app.route('/get_continuation', methods=["POST"])
@login_required
def get_continuation(username: str):
    return ai_assist(request.form.get("cont"), "续写")


def ai_assist(text: str, action: str) -> str | Response:
    try:
        key = dbhelper.get_token("aistudio")
    except RuntimeError:
        return jsonify({
            "success": False,
            "message": "A token is needed to use AI related features.",
        })

    ask_prompt = f"帮我{action}下面这段话:\n{text}"

    erniebot.api_type = "aistudio"
    erniebot.access_token = key

    try:
        response = erniebot.ChatCompletion.create(
            model='ernie-3.5',
            messages=[{'role': 'user', 'content': ask_prompt}],
        )
        result_text = response['result']
        result_dict = {'answer': result_text}
        print(result_dict)
        return jsonify({
            "success": True,
            "answer": result_text,
        })
    except APIError as e:
        print(f"[!] API error: {e}")
        return jsonify({
            "success": False,
            "message": f"API error: {e}",
        })
    except Exception as e:
        print(f"[!] General exception: {e}")
        return jsonify({
            "success": False,
            "message": f"Exception: {e}",
        })


if __name__ == "__main__":
    app.run(host=HOST, port=PORT, debug=DEBUG)
